package me.flyray.bsin.server;

import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.ComponentScan;
import org.springframework.context.annotation.ImportResource;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.retriever.EmbeddingStoreRetriever;
import dev.langchain4j.retriever.Retriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;

/**
 * @author ：leonard
 * @date ：Created in 2023/04/21 16:00
 * @description：bsin ai 引擎
 * @modified By：
 */

@ImportResource({"classpath*:sofa/rpc-provider-ai.xml"})
@SpringBootApplication
@ComponentScan("me.flyray.bsin.*")
public class BsinAiApplication {

    @Bean
    Retriever<TextSegment> retriever(EmbeddingStore<TextSegment> embeddingStore, EmbeddingModel embeddingModel) {

        // You will need to adjust these parameters to find the optimal setting, which will depend on two main factors:
        // - The nature of your data
        // - The embedding model you are using
        int maxResultsRetrieved = 1;
        double minScore = 0.6;

        return EmbeddingStoreRetriever.from(embeddingStore, embeddingModel, maxResultsRetrieved, minScore);
    }

    @Bean
    EmbeddingModel embeddingModel() {
        // 此处可以选择不同的EmbeddingModel：OpenAiEmbeddingModel
        return new AllMiniLmL6V2EmbeddingModel();
    }

    @Bean
    EmbeddingStore<TextSegment> embeddingStore(){

        // Normally, you would already have your embedding store filled with your data.
        // However, for the purpose of this demonstration, we will:
        EmbeddingStore<TextSegment> embeddingStore = MilvusEmbeddingStore.builder()
                .host("localhost")
                .port(19530)
                .collectionName("bolei_test")
                .dimension(384)
                .build();
        // 1. Create an in-memory embedding store
        // EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        return embeddingStore;
    }

    public static void main(String[] args) {
        SpringApplication springApplication = new SpringApplication(BsinAiApplication.class);
        springApplication.run(args);
    }

}
